import math


def euclidean_distance(x1, x2):
    if len(x1) != len(x2):
        raise ValueError("两个样本必须具有相同的维度")
    # 欧氏距离公式：sqrt(sum((x1_i - x2_i)^2))
    return math.sqrt(sum((a - b) ** 2 for a, b in zip(x1, x2)))


def knn_classify(train_features, train_labels, test_sample, k):
    # 输入合法性检查
    if not train_features:
        raise ValueError("训练样本不能为空")
    if len(train_features) != len(train_labels):
        raise ValueError("训练特征与标签数量必须一致")
    if k <= 0:
        raise ValueError("k必须为正整数")
    # 若k大于训练样本数，强制使用所有样本
    if k > len(train_features):
        k = len(train_features)
        print(f"警告：k值大于训练样本数，已自动调整为{k}")
    # 检查特征维度是否一致
    feature_dim = len(train_features[0])
    if len(test_sample) != feature_dim:
        raise ValueError(f"测试样本维度（{len(test_sample)}）与训练样本维度（{feature_dim}）不一致")

    # 1. 计算测试样本与所有训练样本的距离
    distances = []
    for i in range(len(train_features)):
        dist = euclidean_distance(test_sample, train_features[i])
        distances.append((dist, train_labels[i]))  # 存储（距离，标签）元组

    # 2. 按距离升序排序（距离越小越近）
    distances.sort(key=lambda x: x[0])  # 按元组第一个元素（距离）排序

    # 3. 取前k个最近邻的标签
    k_nearest_labels = [item[1] for item in distances[:k]]

    # 4. 投票：选择出现次数最多的标签作为预测结果
    label_count = {}
    for label in k_nearest_labels:
        if label in label_count:
            label_count[label] += 1
        else:
            label_count[label] = 1

    # 处理票数相同的情况（取第一个出现的最大票数标签）
    max_count = -1
    predicted_label = None
    for label, count in label_count.items():
        if count > max_count:
            max_count = count
            predicted_label = label

    return predicted_label


def knn_predict_batch(train_features, train_labels, test_samples, k):
    return [knn_classify(train_features, train_labels, sample, k) for sample in test_samples]

# 训练数据（二维特征，标签为0或1）
train_features = [
    [1.2, 2.3], [1.8, 1.9], [2.1, 2.8], [2.5, 2.2],  # 标签0
    [5.3, 6.1], [6.0, 5.8], [6.2, 6.5], [7.0, 6.3]   # 标签1
]
train_labels = [0, 0, 0, 0, 1, 1, 1, 1]

# 测试样本
test_samples = [
    [1.5, 2.5],  # 预计属于0
    [6.5, 6.0],  # 预计属于1
    [3.0, 3.5]   # 靠近0的样本，预计属于0
]

# 预测（k=3）
predictions = knn_predict_batch(train_features, train_labels, test_samples, k=3)

print("测试样本：", test_samples)
print("预测标签：", predictions)